import bisect
from math import floor, ceil
from itertools import chain
import numpy as np
import pandas as pd

import matplotlib.dates as dt


def reject_outliers(data, whitelist_range, m=0.5):
  d = np.abs(data - np.median(data))
  mdev = np.median(d)
  s = d / mdev if mdev else 0.
  whitelist_ind = (whitelist_range[0] <= data) & (data <= whitelist_range[1])
  return data[(s < m) | whitelist_ind]


def unzip_tpq(tpq, tss):
  tsso = []
  prso = []
  qtyso = []
  for ts, pr, qty in tpq:
    if tss[0] < ts < tss[-1]:
      tsso.append(ts)
      prso.append(pr)
      qtyso.append(qty)
  return tsso, prso, qtyso


def to_updates_and_plot_bbo(ax,
                            tss,
                            obs,
                            ask0s,
                            bid0s,
                            trd_tpq,
                            self_tss,
                            self_obs,
                            fill_tpq,
                            sub_tpq,
                            acc_tpq,
                            can_tpq):
  tick_bins_cnt = 2000
  step_bins_cnt = 2000

  tssfill, prcfill, qtyfill = unzip_tpq(fill_tpq, tss)
  tsssub, prcsub, qtysub = unzip_tpq(sub_tpq, tss)
  tssacc, prcacc, qtyacc = unzip_tpq(acc_tpq, tss)
  tsscan, prccan, qtycan = unzip_tpq(can_tpq, tss)
  tsstrd, prctrd, qtytrd = unzip_tpq(trd_tpq, tss)

  prices = []
  for ob in obs:
    prices.extend(ob.keys())
  if prcacc is not None and len(prcacc) > 0:
    prices.extend([min(prcacc), max(prcacc)])
  prices = np.unique(prices)

  rejected = reject_outliers(np.array(prices),
                             (min(chain(prcacc, bid0s)), max(chain(prcacc, ask0s))))

  sizes, boundaries = np.histogram(rejected, tick_bins_cnt)

  def into_tick_bins(prices):
    inds = np.searchsorted(boundaries, prices, side='left')
    inds2 = inds.copy()
    inds2[inds2 >= len(boundaries)] = 0

    ret = boundaries[inds2]
    ret[inds >= len(boundaries)] = np.nan
    ret[inds == 0] = np.nan
    return ret

  def into_tick_bin(price):
    # ind = bisect.bisect_left(boundaries, price)
    ind = np.searchsorted(boundaries, price, side='left')
    if ind >= len(boundaries):
      return np.nan
    if ind == 0:
      return np.nan
    return boundaries[ind]
    # for (s, b) in zip(boundaries, boundaries[1:]):
    #  if b > price > s:
    #    return s
    # return False

  min_ts = tss[0]
  max_ts = tss[-1]
  step_thresholds = range(int(floor(min_ts)),
                          int(ceil(max_ts)),
                          int(floor((max_ts - min_ts) / (step_bins_cnt))))
  step_thresholds_ts = np.array(pd.DatetimeIndex(step_thresholds).tolist())

  def into_step_bins(times):
    inds = np.searchsorted(step_thresholds, times, side='left') + 1
    inds2 = inds.copy()
    inds2[inds2 >= len(step_thresholds)] = 0
    ret = step_thresholds_ts[inds2]
    ret[inds >= len(step_thresholds)] = False
    return ret

  def into_step_bin(time):
    ind = bisect.bisect_left(step_thresholds, time) + 1
    if ind >= len(step_thresholds):
      return False
    return step_thresholds_ts[ind]
    # for (s, b) in zip(step_thresholds, step_thresholds[1:]):
    #  if b > time > s:
    #    return b
    # return False

  # tssbin = np.array([into_step_bin(ts) for ts in tss])
  # askbin = np.array([into_tick_bin(ask0) for ask0 in ask0s])
  # bidbin = np.array([into_tick_bin(bid0) for bid0 in bid0s])
  tssbin = into_step_bins(tss)
  askbin = into_tick_bins(ask0s)
  bidbin = into_tick_bins(bid0s)

  askbin = askbin[tssbin != False]
  bidbin = bidbin[tssbin != False]
  tssbin = tssbin[tssbin != False]

  xfmt = dt.DateFormatter('%H:%M:%S.%f')
  ax.xaxis.set_major_formatter(xfmt)

  ax.plot(tssbin, askbin, lw=1, color='red', ls='--', drawstyle='steps-post', alpha=0.5)
  ax.plot(tssbin, bidbin, lw=1, color='green', ls='--', drawstyle='steps-post', alpha=0.5)

  def to_bin(tssin, prcin, qtyin):
    tssbin = []
    prcbin = []
    qtyout = []

    tssbins = into_step_bins(tssin)
    prcbins = into_tick_bins(prcin)

    for tb, ts, pr, pb, qty in zip(tssbins, tssin, prcin, prcbins, qtyin):
      if tb != False:
        tssbin.append(tb)
        prcbin.append(pb)
        qtyout.append(qty)
    return np.array(tssbin), np.array(prcbin), np.array(qtyout)

  tssfillbin, prcfillbin, qtyfill = to_bin(tssfill, prcfill, qtyfill)

  tsssubbin, prcsubbin, qtysub = to_bin(tsssub, prcsub, qtysub)
  ax.scatter(tsssubbin[qtysub > 0],
             prcsubbin[qtysub > 0],
             color='green',
             facecolors='none',
             linestyle=':',
             alpha=0.5)
  ax.scatter(tsssubbin[qtysub < 0],
             prcsubbin[qtysub < 0],
             color='red',
             facecolors='none',
             linestyle=':',
             alpha=0.5)

  ax.scatter(tssfillbin[qtyfill > 0], prcfillbin[qtyfill > 0], color='green', s=25)
  ax.scatter(tssfillbin[qtyfill < 0], prcfillbin[qtyfill < 0], color='red', s=25)

  tssaccbin, prcaccbin, qtyacc = to_bin(tssacc, prcacc, qtyacc)
  '''
  ax.scatter(tssaccbin[qtyacc > 0], prcaccbin[qtyacc > 0], color='green', facecolors='none', s=15)
  ax.scatter(tssaccbin[qtyacc < 0], prcaccbin[qtyacc < 0], color='red', facecolors='none', s=15)'''

  tsscanbin, prccanbin, qtycan = to_bin(tsscan, prccan, qtycan)
  ax.scatter(tsscanbin[qtycan > 0], prccanbin[qtycan > 0], color='green', marker='|', s=15)
  ax.scatter(tsscanbin[qtycan < 0], prccanbin[qtycan < 0], color='red', marker='|', s=15)

  tsstrdbin, prctrdbin, qtytrd = to_bin(tsstrd, prctrd, qtytrd)
  ax.scatter(tsstrdbin[qtytrd > 0],
             prctrdbin[qtytrd > 0],
             color='green',
             marker='*',
             s=15,
             alpha=0.5)
  ax.scatter(tsstrdbin[qtytrd < 0], prctrdbin[qtytrd < 0], color='red', marker='*', s=15, alpha=0.5)

  # for debug
  # ax.vlines(tssbin[:-1], ymin=min(bidbin), ymax=max(askbin), color='yellow', lw=0.5)

  def obs_to_updates(tss_in, obs_in):
    updates = {}

    tickcache = {}

    for i, (ts, orderbook) in enumerate(zip(tss_in, obs_in)):
      next_ts = tss_in[i + 1] if i + 1 < len(tss_in) else tss_in[i]
      time = into_step_bin(ts)
      next_time = into_step_bin(next_ts)
      for price, size in orderbook.items():
        if price not in tickcache:
          tickcache[price] = into_tick_bin(price)
        price = tickcache[price]
        if not price or not time:
          continue
        if price not in updates:
          updates[price] = []
        if time and next_time:
          updates[price].append((time, next_time, size))
    return updates

  last_ind = bisect.bisect_left(self_tss, tss[-1])

  updates = obs_to_updates(tss, obs)

  self_tss = self_tss[:last_ind]
  self_obs = self_obs[:last_ind]

  self_updates = obs_to_updates(self_tss, self_obs)

  return updates, self_updates


def colorfunc_blue(max_threshold, min_threshold, sdet):
  if sdet < min_threshold:
    return (0, 0, 0)
  elif sdet > max_threshold:
    return (0, 1, 1)
  else:
    return (0, sdet / max_threshold, sdet / max_threshold)


def get_colorfunc_blue(max_threshold, min_threshold):
  return lambda sdet: colorfunc_blue(max_threshold, min_threshold, sdet)


def get_hlines_from_vtuples(price, vtuples, colorfunc):
  ys = []
  xmins = []
  xmaxs = []
  colors = []
  for t1, t2, sdet in vtuples:
    xmins.append(t1)
    xmaxs.append(t2)
    ys.append(price)
    # for debug
    # colors.append((0, 1, 1));continue
    colors.append(colorfunc(sdet))

  # revive this line for slow mode.
  # return ys, xmins, xmaxs, colors

  tups = list(zip(ys, xmins, xmaxs, colors))
  tups.append(None)
  # squeeze hlines for speedup
  tups_sqz = []
  prev_tup = list(tups[0])
  for tup in tups[1:]:
    if tup and prev_tup[2] == tup[1] and prev_tup[3] == tup[3]:
      prev_tup[2] = tup[2]
    else:
      tups_sqz.append(prev_tup)
      if tup is not None:
        prev_tup = list(tup)

  return zip(*tups_sqz)


def plot_price_levels(ax, updates, colorfunc, lw=0.5, zorder=0, linestyles='solid'):
  ys = []
  xmins = []
  xmaxs = []
  colors = []

  for price, vtuples in updates.items():
    if len(vtuples) == 0:
      continue
    ysp, xminsp, xmaxsp, colorsp = get_hlines_from_vtuples(price, vtuples, colorfunc)
    ys.extend(ysp)
    xmins.extend(xminsp)
    xmaxs.extend(xmaxsp)
    colors.extend(colorsp)

  ax.xaxis_date()
  for tick in ax.get_xticklabels():
    tick.set_rotation(20)
  ax.hlines(ys, xmins, xmaxs, color=colors, lw=lw, alpha=1, zorder=zorder, linestyles=linestyles)


def plot_cross(ax, alongts, crsbook):
  df = pd.DataFrame(crsbook)
  df = df.loc[(alongts[0] <= df['timestamp']) & (df['timestamp'] <= alongts[-1])]
  pdi = pd.DatetimeIndex(df['timestamp'])
  # ax.twinx()
  ax.plot(pdi, df['bid0p'], lw=1, color='gray', ls='--', drawstyle='steps-post')
  ax.plot(pdi, df['ask0p'], lw=1, color='gray', ls='--', drawstyle='steps-post')
